# -*- coding: utf-8 -*-
"""
Created on Sat Apr  9 09:35:19 2022

@author: S1mple
"""

#import os
import csv



def read_split_data(root: str):
    

    train_images_path = []  # 存储训练集的所有图片路径
    train_images_labels = []  # 存储训练集图片对应索引信息
    test_images_path = []  # 存储验证集的所有图片路径
    test_images_labels = []  # 存储验证集图片对应索引信息
    val_images_path = []  # 存储验证集的所有图片路径
    val_images_labels = []  # 存储验证集图片对应索引信息
    # 遍历每个文件夹下的文件
    root = "./data"


    reader = csv.reader(open(root+"/train.csv",'r',encoding='utf-8',newline=''))
    next(reader)
    labels_count = -1
    labels_pre = ""
    for filename,labels in reader: 
        train_images_path.append(root+"/image/images/"+filename)
        if labels != labels_pre:
            labels_pre = labels
            labels_count += 1
            train_images_labels.append(labels_count)
        else :
            train_images_labels.append(labels_count)
            
        
    reader = csv.reader(open(root+"/test.csv",'r',encoding='utf-8',newline=''))
    next(reader)
    labels_count = -1
    labels_pre = ""
    for filename,labels in reader: 
        test_images_path.append(root+"/image/images/"+filename)
        if labels != labels_pre:
            labels_pre = labels
            labels_count += 1
            test_images_labels.append(labels_count)
        else :
            test_images_labels.append(labels_count)
            
    reader = csv.reader(open(root+"/val.csv",'r',encoding='utf-8',newline=''))
    next(reader)
    labels_count = -1
    labels_pre = ""
    for filename,labels in reader: 
        val_images_path.append(root+"/image/images/"+filename)
        if labels != labels_pre:
            labels_pre = labels
            labels_count += 1
            val_images_labels.append(labels_count)
        else :
           val_images_labels.append(labels_count)
    print("{} images for training.".format(len(train_images_path)))
    print("{} images for test.".format(len(test_images_path)))
    print("{} images for val.".format(len(val_images_path)))

 

    return train_images_path, train_images_labels, test_images_path, test_images_labels, val_images_path, val_images_labels
